{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "17da01e4",
   "metadata": {},
   "source": [
    "计算Q时去除基线,基线使用baseline模型估计\n",
    "\n",
    "Q是用蒙特卡洛法估计的,baseline模型的loss就是它的计算结果和Q求mse loss\n",
    "\n",
    "Q在使用时减去baseline模型的计算结果,相当于去基线"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "91162df8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/P9b71AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAUjklEQVR4nO3df2xT570G8Md2bCeQHIeQxW5ELJDWjmb86BYgeExa1XqkLJrGGt39EOqyClGVOVxoNqRGaumlt1I6Jt1u3Wj4oxrwD2PKJDoRUaootOF2dUkJixQC5K677U0uxXYh8nESiO3Y3/vHbk4xCSwhiV/beT7SkfD7vra/54Afjs/rc45JRARERAqYVRdARAsXA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRRFkAHDx7E8uXLkZ+fj+rqanR1dakqhYgUURJAf/zjH9HY2IiXXnoJFy5cwNq1a1FTU4NQKKSiHCJSxKTiZNTq6mqsX78ev/vd7wAAyWQSFRUV2LVrF55//vl0l0NEiuSl+w1jsRi6u7vR1NRktJnNZni9Xvj9/imfE41GEY1GjcfJZBJDQ0NYunQpTCbTvNdMRDMjIhgeHkZ5eTnM5rt/0Up7AF2/fh2JRAJOpzOl3el04sqVK1M+p7m5Gfv3709HeUQ0hwYHB7Fs2bK79qc9gO5HU1MTGhsbjce6rsPtdmNwcBCapimsjIimEolEUFFRgaKionuOS3sAlZaWwmKxIBgMprQHg0G4XK4pn2O322G32ye1a5rGACLKYP/sEEnaZ8FsNhuqqqrQ0dFhtCWTSXR0dMDj8aS7HCJSSMlXsMbGRtTX12PdunXYsGEDfv3rX2N0dBRPP/20inKISBElAfTDH/4Qn3/+Ofbt24dAIIBHHnkEp0+fnnRgmohym5LfAc1WJBKBw+GArus8BkSUgab7GeW5YESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlJmxgF09uxZfPe730V5eTlMJhPeeuutlH4Rwb59+/DAAw+goKAAXq8Xf/vb31LGDA0NYdu2bdA0DcXFxdi+fTtGRkZmtSJElH1mHECjo6NYu3YtDh48OGX/gQMH8Prrr+PQoUM4d+4cFi9ejJqaGoyNjRljtm3bhr6+PrS3t6OtrQ1nz57FM888c/9rQUTZSWYBgJw4ccJ4nEwmxeVyya9+9SujLRwOi91ulz/84Q8iInLp0iUBIB999JEx5u233xaTySRXr16d1vvqui4ARNf12ZRPRPNkup/ROT0G9MknnyAQCMDr9RptDocD1dXV8Pv9AAC/34/i4mKsW7fOGOP1emE2m3Hu3LkpXzcajSISiaQsRJT95jSAAoEAAMDpdKa0O51Ooy8QCKCsrCylPy8vDyUlJcaYOzU3N8PhcBhLRUXFXJZNRIpkxSxYU1MTdF03lsHBQdUlEdEcmNMAcrlcAIBgMJjSHgwGjT6Xy4VQKJTSPz4+jqGhIWPMnex2OzRNS1mIKPvNaQCtWLECLpcLHR0dRlskEsG5c+fg8XgAAB6PB+FwGN3d3caYM2fOIJlMorq6ei7LIaIMlzfTJ4yMjODjjz82Hn/yySfo6elBSUkJ3G439uzZg1deeQUPPvggVqxYgRdffBHl5eXYunUrAODhhx/GE088gR07duDQoUOIx+NoaGjAj370I5SXl8/ZihFRFpjp9Nq7774rACYt9fX1IvKPqfgXX3xRnE6n2O12efzxx6W/vz/lNW7cuCE//vGPpbCwUDRNk6efflqGh4fnfIqPiNSY7mfUJCKiMP/uSyQSgcPhgK7rPB5ElIGm+xnNilkwIspNDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhImRnflodovogIbl7/H4yFv7ixpXVxMYoeeAgmk0lhZTRfGECUUT6//D4+v3zWeOxwr0bRAw8prIjmE7+CUeYQQXI8proKSiMGEGUMkSSS41HVZVAaMYAoc4hgPHYrpcls4VGCXMYAooyRTMQRve0ANAAULK1QVA2lAwOIMsqddwq32AoUVULpwACijGax2lWXQPNoRgHU3NyM9evXo6ioCGVlZdi6dSv6+/tTxoyNjcHn82Hp0qUoLCxEXV0dgsHU3eqBgQHU1tZi0aJFKCsrw969ezE+Pj77taGcY7Hmqy6B5tGMAqizsxM+nw8ffvgh2tvbEY/HsXnzZoyOjhpjnnvuOZw8eRKtra3o7OzEZ599hieffNLoTyQSqK2tRSwWwwcffICjR4/iyJEj2Ldv39ytFWUlSSYApH4FM5nN/BFiLpNZCIVCAkA6OztFRCQcDovVapXW1lZjzOXLlwWA+P1+ERE5deqUmM1mCQQCxpiWlhbRNE2i0ei03lfXdQEguq7PpnzKMGORG3LhSKN0HdphLEP/3a26LLoP0/2MzuoYkK7rAICSkhIAQHd3N+LxOLxerzFm5cqVcLvd8Pv9AAC/34/Vq1fD6XQaY2pqahCJRNDX1zfl+0SjUUQikZSFck8yEYNIUnUZlEb3HUDJZBJ79uzBpk2bsGrVKgBAIBCAzWZDcXFxylin04lAIGCMuT18Jvon+qbS3NwMh8NhLBUVnJrNRcnx2P9/DbuNifMkuey+/3Z9Ph8uXryI48ePz2U9U2pqaoKu68YyODg47+9J6ReNfI5kfMx4bLEVwF5UqrAimm/39TPThoYGtLW14ezZs1i2bJnR7nK5EIvFEA6HU/aCgsEgXC6XMaarqyvl9SZmySbG3Mlut8Nu53RsrpNk6tcvkzkPZk7D57QZ7QGJCBoaGnDixAmcOXMGK1asSOmvqqqC1WpFR0eH0dbf34+BgQF4PB4AgMfjQW9vL0KhkDGmvb0dmqahsrJyNutCOcZkNsNssaoug+bRjPaAfD4fjh07hj//+c8oKioyjtk4HA4UFBTA4XBg+/btaGxsRElJCTRNw65du+DxeLBx40YAwObNm1FZWYmnnnoKBw4cQCAQwAsvvACfz8e9nAVM7vgFNACYTGaY8xhAuWxGAdTS0gIAePTRR1PaDx8+jJ/+9KcAgNdeew1msxl1dXWIRqOoqanBG2+8YYy1WCxoa2vDzp074fF4sHjxYtTX1+Pll1+e3ZpQ1pPkHT9GNZlgMlnUFENpYZKp/uvJcJFIBA6HA7quQ9M01eXQHBARhC51YuD9Y0abrbAEq/7l32Cx8dfQ2Wa6n1HOcVLGuH0GjBYGBhBljET8zouR8RSMXMcAooxxa+hqyuP8YidMFh4DymUMIMoYd14POi+/ECb+Ejqn8W+XMpY5zwZ+DcttDCDKWJY8O8BLceQ0BhBlBpHJp2Lw+E/OYwBRRhBJQBLxSe28GFluYwBRRpBkAskEL8u70DCAKCMkE+OTbkpoMvOeYLmOAUQZITE2itjI0BcNJhMKSsrVFURpwQCijCCQO86IN/GOGAsAA4gyFi9GlvsYQJSRTCYTb0q4ADCAKCNMfU8wHoTOdQwgygjJ8did+QOehpH7GECUEZLxKKZIIMpxDCDKCInY2KRZMO4A5T4GEGWEW0NXgdvuimpdXAxrfqHCiigdGECUEZJ3XJDenGeDibfkyXkMIMpIZkseTGaeDZ/rGECUkUxmBtBCwAAi5UQk5fgPAJjMFgbQAsAAooyQiMf++SDKOfypKaXF6Ogo4vHJFxwDAJEkxm4Op7QlEuPQw+G7XpK1qKgIFl4xMesxgCgtfvGLX+DkyZNT9lnMJuzeuhbfrPzi8hvnz5/Hln9dPeVPE202G06fPo2HHnponqqldGEAUVoMDQ3h6tWrU/bZrRYsKXoUfaPfxM2EBnf+JVz+9DL+9y7jbTbbXfemKLvM6BhQS0sL1qxZA03ToGkaPB4P3n77baN/bGwMPp8PS5cuRWFhIerq6hAMBlNeY2BgALW1tVi0aBHKysqwd+9ejI/zUpwLmcCCj2OPYXDsYdyIL0PvyKP4eGiJ6rIoDWYUQMuWLcOrr76K7u5unD9/Ho899hi+973voa+vDwDw3HPP4eTJk2htbUVnZyc+++wzPPnkk8bzE4kEamtrEYvF8MEHH+Do0aM4cuQI9u3bN7drRVlFYMJIohgT516Miw1DN3kpjgVBZmnJkiXy5ptvSjgcFqvVKq2trUbf5cuXBYD4/X4RETl16pSYzWYJBALGmJaWFtE0TaLR6LTfU9d1ASC6rs+2fEqTH/zgB4J/nG06abFbrfLL538p//6KX/a/0iX/0Xxanti08a7jbTabXLx4UfUq0T1M9zN638eAEokEWltbMTo6Co/Hg+7ubsTjcXi9XmPMypUr4Xa74ff7sXHjRvj9fqxevRpOp9MYU1NTg507d6Kvrw9f+9rXZlTDlStXUFjI84WyQSQSuWvfeGIc/r8cxVjeX3ArWYgy2wA+/vS/7jpeRPD3v/+dt+zJYCMjI9MaN+MA6u3thcfjwdjYGAoLC3HixAlUVlaip6cHNpsNxcXFKeOdTicCgQAAIBAIpITPRP9E391Eo1FEo1/cMWHiH7Ou6zx+lCVisbv/zieRFLz1n5cAXJr26w0PDyMcDs++MJoXo6Oj0xo34wD6yle+gp6eHui6jj/96U+or69HZ2fnjAuciebmZuzfv39Se3V1NTRNm9f3prlRWlo6Z69lMpnwyCOP4Ktf/eqcvSbNrXvt8d5uxr+Ettls+PKXv4yqqio0Nzdj7dq1+M1vfgOXy4VYLDbpf6VgMAiXywUAcLlck2bFJh5PjJlKU1MTdF03lsHBwZmWTUQZaNanYiSTSUSjUVRVVcFqtaKjo8Po6+/vx8DAADweDwDA4/Ggt7cXoVDIGNPe3g5N01BZWXnX97Db7cbU/8RCRNlvRl/BmpqasGXLFrjdbgwPD+PYsWN477338M4778DhcGD79u1obGxESUkJNE3Drl274PF4sHHjRgDA5s2bUVlZiaeeegoHDhxAIBDACy+8AJ/PB7ud065EC82MAigUCuEnP/kJrl27BofDgTVr1uCdd97Bt7/9bQDAa6+9BrPZjLq6OkSjUdTU1OCNN94wnm+xWNDW1oadO3fC4/Fg8eLFqK+vx8svvzy3a0UZp6CgYM72XG02G8xmnkedC0wiknVXAo9EInA4HNB1nV/HssT169dx69atOXs9l8sFq5VXTMxU0/2M8lwwSou5nAWj3MH9WCJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKRMnuoC7oeIAAAikYjiSohoKhOfzYnP6t1kZQDduHEDAFBRUaG4EiK6l+HhYTgcjrv2Z2UAlZSUAAAGBgbuuXKUKhKJoKKiAoODg9A0TXU5WYHb7P6ICIaHh1FeXn7PcVkZQGbzPw5dORwO/qO4D5qmcbvNELfZzE1n54AHoYlIGQYQESmTlQFkt9vx0ksvwW63qy4lq3C7zRy32fwyyT+bJyMimidZuQdERLmBAUREyjCAiEgZBhARKZOVAXTw4EEsX74c+fn5qK6uRldXl+qSlGlubsb69etRVFSEsrIybN26Ff39/SljxsbG4PP5sHTpUhQWFqKurg7BYDBlzMDAAGpra7Fo0SKUlZVh7969GB8fT+eqKPPqq6/CZDJhz549Rhu3WZpIljl+/LjYbDb5/e9/L319fbJjxw4pLi6WYDCoujQlampq5PDhw3Lx4kXp6emR73znO+J2u2VkZMQY8+yzz0pFRYV0dHTI+fPnZePGjfKNb3zD6B8fH5dVq1aJ1+uVv/71r3Lq1CkpLS2VpqYmFauUVl1dXbJ8+XJZs2aN7N6922jnNkuPrAugDRs2iM/nMx4nEgkpLy+X5uZmhVVljlAoJACks7NTRETC4bBYrVZpbW01xly+fFkAiN/vFxGRU6dOidlslkAgYIxpaWkRTdMkGo2mdwXSaHh4WB588EFpb2+Xb33rW0YAcZulT1Z9BYvFYuju7obX6zXazGYzvF4v/H6/wsoyh67rAL44Ybe7uxvxeDxlm61cuRJut9vYZn6/H6tXr4bT6TTG1NTUIBKJoK+vL43Vp5fP50NtbW3KtgG4zdIpq05GvX79OhKJRMpfOgA4nU5cuXJFUVWZI5lMYs+ePdi0aRNWrVoFAAgEArDZbCguLk4Z63Q6EQgEjDFTbdOJvlx0/PhxXLhwAR999NGkPm6z9MmqAKJ78/l8uHjxIt5//33VpWS0wcFB7N69G+3t7cjPz1ddzoKWVV/BSktLYbFYJs1GBINBuFwuRVVlhoaGBrS1teHdd9/FsmXLjHaXy4VYLIZwOJwy/vZt5nK5ptymE325pru7G6FQCF//+teRl5eHvLw8dHZ24vXXX0deXh6cTie3WZpkVQDZbDZUVVWho6PDaEsmk+jo6IDH41FYmToigoaGBpw4cQJnzpzBihUrUvqrqqpgtVpTtll/fz8GBgaMbebxeNDb24tQKGSMaW9vh6ZpqKysTM+KpNHjjz+O3t5e9PT0GMu6deuwbds248/cZmmi+ij4TB0/flzsdrscOXJELl26JM8884wUFxenzEYsJDt37hSHwyHvvfeeXLt2zVhu3rxpjHn22WfF7XbLmTNn5Pz58+LxeMTj8Rj9E1PKmzdvlp6eHjl9+rR86UtfWlBTyrfPgolwm6VL1gWQiMhvf/tbcbvdYrPZZMOGDfLhhx+qLkkZAFMuhw8fNsbcunVLfvazn8mSJUtk0aJF8v3vf1+uXbuW8jqffvqpbNmyRQoKCqS0tFR+/vOfSzweT/PaqHNnAHGbpQcvx0FEymTVMSAiyi0MICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISJn/A+Gay79kDxcKAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import gym\n",
    "\n",
    "\n",
    "#定义环境\n",
    "class MyWrapper(gym.Wrapper):\n",
    "\n",
    "    def __init__(self):\n",
    "        env = gym.make('CartPole-v1', render_mode='rgb_array')\n",
    "        super().__init__(env)\n",
    "        self.env = env\n",
    "        self.step_n = 0\n",
    "\n",
    "    def reset(self):\n",
    "        state, _ = self.env.reset()\n",
    "        self.step_n = 0\n",
    "        return state\n",
    "\n",
    "    def step(self, action):\n",
    "        state, reward, terminated, truncated, info = self.env.step(action)\n",
    "        over = terminated or truncated\n",
    "\n",
    "        #限制最大步数\n",
    "        self.step_n += 1\n",
    "        if self.step_n >= 200:\n",
    "            over = True\n",
    "        \n",
    "        #没坚持到最后,扣分\n",
    "        if over and self.step_n < 200:\n",
    "            reward = -1000\n",
    "\n",
    "        return state, reward, over\n",
    "\n",
    "    #打印游戏图像\n",
    "    def show(self):\n",
    "        from matplotlib import pyplot as plt\n",
    "        plt.figure(figsize=(3, 3))\n",
    "        plt.imshow(self.env.render())\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "env = MyWrapper()\n",
    "\n",
    "env.reset()\n",
    "\n",
    "env.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ecfbe912",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Sequential(\n",
       "   (0): Linear(in_features=4, out_features=64, bias=True)\n",
       "   (1): ReLU()\n",
       "   (2): Linear(in_features=64, out_features=64, bias=True)\n",
       "   (3): ReLU()\n",
       "   (4): Linear(in_features=64, out_features=2, bias=True)\n",
       "   (5): Softmax(dim=1)\n",
       " ),\n",
       " Sequential(\n",
       "   (0): Linear(in_features=4, out_features=64, bias=True)\n",
       "   (1): ReLU()\n",
       "   (2): Linear(in_features=64, out_features=64, bias=True)\n",
       "   (3): ReLU()\n",
       "   (4): Linear(in_features=64, out_features=1, bias=True)\n",
       " ))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "#定义模型,计算每个动作的概率\n",
    "model_action = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 2),\n",
    "    torch.nn.Softmax(dim=1),\n",
    ")\n",
    "\n",
    "#基线模型,评估state的价值\n",
    "model_baseline = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 1),\n",
    ")\n",
    "\n",
    "model_action, model_baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "84cbf0ff",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_7326/3856618961.py:30: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at  ../torch/csrc/utils/tensor_new.cpp:201.)\n",
      "  state = torch.FloatTensor(state).reshape(-1, 4)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "-989.0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython import display\n",
    "import random\n",
    "\n",
    "\n",
    "#玩一局游戏并记录数据\n",
    "def play(show=False):\n",
    "    state = []\n",
    "    action = []\n",
    "    reward = []\n",
    "\n",
    "    s = env.reset()\n",
    "    o = False\n",
    "    while not o:\n",
    "        #根据概率采样\n",
    "        prob = model_action(torch.FloatTensor(s).reshape(1, 4))[0].tolist()\n",
    "        a = random.choices(range(2), weights=prob, k=1)[0]\n",
    "\n",
    "        ns, r, o = env.step(a)\n",
    "\n",
    "        state.append(s)\n",
    "        action.append(a)\n",
    "        reward.append(r)\n",
    "\n",
    "        s = ns\n",
    "\n",
    "        if show:\n",
    "            display.clear_output(wait=True)\n",
    "            env.show()\n",
    "\n",
    "    state = torch.FloatTensor(state).reshape(-1, 4)\n",
    "    action = torch.LongTensor(action).reshape(-1, 1)\n",
    "    reward = torch.FloatTensor(reward).reshape(-1, 1)\n",
    "\n",
    "    return state, action, reward, reward.sum().item()\n",
    "\n",
    "\n",
    "state, action, reward, reward_sum = play()\n",
    "\n",
    "reward_sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "14ac0e81",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer_action = torch.optim.Adam(model_action.parameters(), lr=5e-3)\n",
    "optimizer_baseline = torch.optim.Adam(model_baseline.parameters(), lr=5e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0342b2ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([12, 1])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_value(reward):\n",
    "    #计算当前state的价值,其实就是Q(state,action),这里是用蒙特卡洛法估计的\n",
    "    value = []\n",
    "    for i in range(len(reward)):\n",
    "        s = 0\n",
    "        for j in range(i, len(reward)):\n",
    "            s += reward[j] * 0.99**(j - i)\n",
    "        value.append(s)\n",
    "\n",
    "    return torch.FloatTensor(value).reshape(-1, 1)\n",
    "\n",
    "\n",
    "value = get_value(reward)\n",
    "\n",
    "value.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "99cc52b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([12, 1])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#训练baseline模型\n",
    "def train_baseline(state, value):\n",
    "    baseline = model_baseline(state)\n",
    "\n",
    "    loss = torch.nn.functional.mse_loss(baseline, value)\n",
    "    loss.backward()\n",
    "    optimizer_baseline.step()\n",
    "    optimizer_baseline.zero_grad()\n",
    "\n",
    "    return baseline.detach()\n",
    "\n",
    "\n",
    "baseline = train_baseline(state, value)\n",
    "\n",
    "baseline.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "42024b65",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-600.0712280273438"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#训练action模型\n",
    "def train_action(state, action, value, baseline):\n",
    "    #重新计算动作的概率\n",
    "    prob = model_action(state).gather(dim=1, index=action)\n",
    "\n",
    "    #求Q最大的导函数 -> partial value / partial action\n",
    "    #注意这里的Q使用前要去基线,这也是baseline模型存在的意义\n",
    "    prob = (prob + 1e-8).log() * (value - baseline)\n",
    "    for i in range(len(prob)):\n",
    "        prob[i] = prob[i] * 0.99**i\n",
    "    loss = -prob.mean()\n",
    "\n",
    "    loss.backward()\n",
    "    optimizer_action.step()\n",
    "    optimizer_action.zero_grad()\n",
    "\n",
    "    return loss.item()\n",
    "\n",
    "\n",
    "train_action(state, action, value, baseline)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "dbe2cd7a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 -655.6979370117188 -981.0\n",
      "100 -32.240211486816406 -773.9\n",
      "200 28.213363647460938 145.75\n",
      "300 31.198749542236328 147.25\n",
      "400 30.996339797973633 200.0\n",
      "500 22.092092514038086 200.0\n",
      "600 24.341480255126953 200.0\n",
      "700 16.99460220336914 149.8\n",
      "800 17.09764862060547 200.0\n",
      "900 18.121004104614258 200.0\n"
     ]
    }
   ],
   "source": [
    "#训练\n",
    "def train():\n",
    "    model_action.train()\n",
    "    model_baseline.train()\n",
    "\n",
    "    #训练N局\n",
    "    for epoch in range(1000):\n",
    "\n",
    "        #一个epoch最少玩N步\n",
    "        steps = 0\n",
    "        while steps < 200:\n",
    "\n",
    "            #玩一局游戏,得到数据\n",
    "            state, action, reward, _ = play()\n",
    "            steps += len(state)\n",
    "\n",
    "            #训练两个模型\n",
    "            value = get_value(reward)\n",
    "            baseline = train_baseline(state, value)\n",
    "            loss = train_action(state, action, value, baseline)\n",
    "\n",
    "        if epoch % 100 == 0:\n",
    "            test_result = sum([play()[-1] for _ in range(20)]) / 20\n",
    "            print(epoch, loss, test_result)\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a101d0e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/P9b71AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAT7klEQVR4nO3db2xT970G8MdOYhOSHKdJFps0sUAqdyziz7oA4YwX3W09sjaaxhpdbRPq0opLL8xBpZmQFqkFgSqFS6XbtRsEadOgbzqmVKIVEW2Va0rQhCGQNloIELX3tkoE2G5AOU5C4yT2977ozdkMgcUk+Jdjno90JHx+X9vfc8h5dHx+/mMTEQERkQJ21Q0Q0cOLAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMooC6ADBw5g8eLFWLBgAaqrq9HZ2amqFSJSREkA/eUvf0FjYyN2796NTz75BKtWrUJNTQ0ikYiKdohIEZuKD6NWV1djzZo1+P3vfw8ASCQSqKiowPbt2/Gb3/wm3e0QkSLZ6X7C8fFxdHV1oampyVxnt9vh8/kQDAanvU8sFkMsFjNvJxIJ3Lx5E8XFxbDZbA+8ZyJKjYhgeHgYZWVlsNvv/kIr7QE0ODiIeDwOt9udtN7tduPKlSvT3qe5uRl79uxJR3tENIcGBgZQXl5+1/G0B9D9aGpqQmNjo3nbMAx4vV4MDAxA0zSFnRHRdKLRKCoqKlBQUHDPurQHUElJCbKyshAOh5PWh8NheDyeae/jdDrhdDrvWK9pGgOIaB77Z5dI0j4L5nA4UFVVhUAgYK5LJBIIBALQdT3d7RCRQkpegjU2NqK+vh6rV6/G2rVr8dvf/hajo6N44YUXVLRDRIooCaCf/exn+Oqrr7Br1y6EQiF897vfxYcffnjHhWkiymxK3gc0W9FoFC6XC4Zh8BoQ0Tw002OUnwUjImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMikH0OnTp/HjH/8YZWVlsNlseO+995LGRQS7du3CokWLkJubC5/Ph88++yyp5ubNm9i0aRM0TUNhYSE2b96MkZGRWW0IEVlPygE0OjqKVatW4cCBA9OO79+/H2+99RYOHTqEc+fOIS8vDzU1NRgbGzNrNm3ahN7eXrS3t6OtrQ2nT5/Giy++eP9bQUTWJLMAQI4dO2beTiQS4vF45PXXXzfXDQ0NidPplD//+c8iInLp0iUBIOfPnzdrPvjgA7HZbHL16tUZPa9hGAJADMOYTftE9IDM9Bid02tAX3zxBUKhEHw+n7nO5XKhuroawWAQABAMBlFYWIjVq1ebNT6fD3a7HefOnZv2cWOxGKLRaNJCRNY3pwEUCoUAAG63O2m92+02x0KhEEpLS5PGs7OzUVRUZNbcrrm5GS6Xy1wqKirmsm0iUsQSs2BNTU0wDMNcBgYGVLdERHNgTgPI4/EAAMLhcNL6cDhsjnk8HkQikaTxyclJ3Lx506y5ndPphKZpSQsRWd+cBtCSJUvg8XgQCATMddFoFOfOnYOu6wAAXdcxNDSErq4us+bkyZNIJBKorq6ey3aIaJ7LTvUOIyMj+Pzzz83bX3zxBbq7u1FUVASv14sdO3bgtddew9KlS7FkyRK8+uqrKCsrw8aNGwEA3/nOd/CjH/0IW7ZswaFDhzAxMYGGhgb8/Oc/R1lZ2ZxtGBFZQKrTax9//LEAuGOpr68XkW+m4l999VVxu93idDrlqaeekr6+vqTHuHHjhvziF7+Q/Px80TRNXnjhBRkeHp7zKT4iUmOmx6hNRERh/t2XaDQKl8sFwzB4PYhoHprpMWqJWTAiykwMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEiZlH+Wh+YHEUH06mVMfj1srsstehS5RY/CZrMp7Ixo5hhAFnatqw0job//Rtuix5/Bo0WPKuyIKDV8CZZBRBKqWyBKCQMokzCAyGIYQBlEEnHVLRClhAGUQSSRwDe/lE1kDQygDMJrQGQ1DKAMIpLgCRBZSkoB1NzcjDVr1qCgoAClpaXYuHEj+vr6kmrGxsbg9/tRXFyM/Px81NXVIRwOJ9X09/ejtrYWCxcuRGlpKXbu3InJycnZb83DLsEzILKWlAKoo6MDfr8fZ8+eRXt7OyYmJrBhwwaMjo6aNS+//DKOHz+O1tZWdHR04Nq1a3j22WfN8Xg8jtraWoyPj+PMmTN4++23ceTIEezatWvutuoh9c1FaJ4CkYXILEQiEQEgHR0dIiIyNDQkOTk50traatZcvnxZAEgwGBQRkRMnTojdbpdQKGTWtLS0iKZpEovFZvS8hmEIADEMYzbtW1oikZBL7/2ndB7aYi6f//cfJBGfVN0a0YyP0VldAzIMAwBQVFQEAOjq6sLExAR8Pp9Zs2zZMni9XgSDQQBAMBjEihUr4Ha7zZqamhpEo1H09vZO+zyxWAzRaDRpoTvxIjRZzX0HUCKRwI4dO7B+/XosX74cABAKheBwOFBYWJhU63a7EQqFzJp/DJ+p8amx6TQ3N8PlcplLRUXF/badUWy25P8+iU/yBRhZyn0HkN/vx8WLF3H06NG57GdaTU1NMAzDXAYGBh74c1pBblFZ0u2xoRDANyOShdzXh1EbGhrQ1taG06dPo7y83Fzv8XgwPj6OoaGhpLOgcDgMj8dj1nR2diY93tQs2VTN7ZxOJ5xO5/20mtHsWY6k28JZMLKYlM6ARAQNDQ04duwYTp48iSVLliSNV1VVIScnB4FAwFzX19eH/v5+6LoOANB1HT09PYhEImZNe3s7NE1DZWXlbLbl4WPn27jI2lI6A/L7/XjnnXfw/vvvo6CgwLxm43K5kJubC5fLhc2bN6OxsRFFRUXQNA3bt2+HrutYt24dAGDDhg2orKzEc889h/379yMUCuGVV16B3+/nWU6KbAwgsriUAqilpQUA8IMf/CBp/eHDh/H8888DAN544w3Y7XbU1dUhFouhpqYGBw8eNGuzsrLQ1taGbdu2Qdd15OXlob6+Hnv37p3dljyEbLYs1S0QzYpNRCw3cRKNRuFyuWAYBjRNU92OEiKC65+ewNXz75vrnFoplv/bLtizHfe4J9GDN9NjlOfwFsYzILI6BpCV8RoQWRz/gi3MZucZEFkbA8jCbn8nNJHV8C/YwjgNT1bHv2AL40VosjoGkEXZbDbgjh8gFFjwXRX0EGMAZRKRbxYii2AAZRCB8DuByFIYQJmEZ0BkMQygDMMzILISBlAGERH+PDNZCgMokwhnwchaGEAZhRehyVoYQBlEeBGaLIYBlFF4BkTWwgDKJDwDIothAFmY7baPYojI//88M5E1MIAszFlQkvT1q/HYLYyPDqlriChFDCALs2Xn3PaBVAGEZ0BkHQwgC/vmC8lu/0Q8kXUwgKyM34hIFse/YAuz2Wx3XIgmshIGkIXxK1nJ6lL6ZVRKLxHB6OgoJicnpx0fHx69420/IyOjkKGhaettNhsKCgpgZ3DRPMEAmueef/55nD17dtoxd2EuXv/39chf8Pep+K1b/wOf/s/gtPWapiEQCGDRokUPpFeiVDGA5jERweDgIK5evTrteOLrfIxM5OF/J9djPLEAS3J7MDh44671IyMjiMc5TU/zR0rn4i0tLVi5ciU0TYOmadB1HR988IE5PjY2Br/fj+LiYuTn56Ourg7hcDjpMfr7+1FbW4uFCxeitLQUO3fuvOtLDLq3WDwb3dF/xbXYUgxOeNE9/CSGJ4tVt0U0YykFUHl5Ofbt24euri5cuHABTz75JH7yk5+gt7cXAPDyyy/j+PHjaG1tRUdHB65du4Znn33WvH88HkdtbS3Gx8dx5swZvP322zhy5Ah27do1t1v1kJhM2DEa1zD1XqDxxALEZKHapohSIbP0yCOPyB//+EcZGhqSnJwcaW1tNccuX74sACQYDIqIyIkTJ8Rut0soFDJrWlpaRNM0icViM35OwzAEgBiGMdv257V4PC5PPPGEAJh2ceXnyX/tPih7Xzsre17rlDf3HZfHl/3L3etdLhkYGFC9WfQQmOkxet/XgOLxOFpbWzE6Ogpd19HV1YWJiQn4fD6zZtmyZfB6vQgGg1i3bh2CwSBWrFgBt9tt1tTU1GDbtm3o7e3F448/nlIPV65cQX5+/v1uwrwnIrh169Zdx8diYzjTcQhRPIYJcaLU8SXCg9fvWp9IJPDZZ58hGo0+iHaJTCMjIzOqSzmAenp6oOs6xsbGkJ+fj2PHjqGyshLd3d1wOBwoLCxMqne73QiFQgCAUCiUFD5T41NjdxOLxRCLxczbUweQYRgZf/3oXtsXm4jj3VN/A/C3GT2WiCAajcLpdM5Rd0TTGx0dnVFdygH07W9/G93d3TAMA++++y7q6+vR0dGRcoOpaG5uxp49e+5YX11dDU3THuhzq5RIJOZ0+7KyslBVVYXy8vI5e0yi6cz0LDvld6Q5HA489thjqKqqQnNzM1atWoU333wTHo8H4+PjGLrtTXDhcBgejwcA4PF47pgVm7o9VTOdpqYmGIZhLgMDA6m2TUTz0KzfEptIJBCLxVBVVYWcnBwEAgFzrK+vD/39/dB1HQCg6zp6enoQiUTMmvb2dmiahsrKyrs+h9PpNKf+pxYisr6UXoI1NTXh6aefhtfrxfDwMN555x2cOnUKH330EVwuFzZv3ozGxkYUFRVB0zRs374duq5j3bp1AIANGzagsrISzz33HPbv349QKIRXXnkFfr+f1yWIHkIpBVAkEsEvf/lLXL9+HS6XCytXrsRHH32EH/7whwCAN954A3a7HXV1dYjFYqipqcHBgwfN+2dlZaGtrQ3btm2DruvIy8tDfX099u7dO7dblUHy8vLm7IxP0zR+ep7mFZuI9b7FPBqNwuVywTCMjH45JiL46quvkmYAZ8Nms8Hj8SA7m5/AoQdrpsco/xLnMZvNhtLSUtVtED0w/F4GIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpEy26gbuh4gAAKLRqOJOiGg6U8fm1LF6N5YMoBs3bgAAKioqFHdCRPcyPDwMl8t113FLBlBRUREAoL+//54bR8mi0SgqKiowMDAATdNUt2MJ3Gf3R0QwPDyMsrKye9ZZMoDs9m8uXblcLv5R3AdN07jfUsR9lrqZnBzwIjQRKcMAIiJlLBlATqcTu3fvhtPpVN2KpXC/pY777MGyyT+bJyMiekAseQZERJmBAUREyjCAiEgZBhARKWPJADpw4AAWL16MBQsWoLq6Gp2dnapbUqa5uRlr1qxBQUEBSktLsXHjRvT19SXVjI2Nwe/3o7i4GPn5+airq0M4HE6q6e/vR21tLRYuXIjS0lLs3LkTk5OT6dwUZfbt2webzYYdO3aY67jP0kQs5ujRo+JwOORPf/qT9Pb2ypYtW6SwsFDC4bDq1pSoqamRw4cPy8WLF6W7u1ueeeYZ8Xq9MjIyYtZs3bpVKioqJBAIyIULF2TdunXy/e9/3xyfnJyU5cuXi8/nk08//VROnDghJSUl0tTUpGKT0qqzs1MWL14sK1eulJdeeslcz32WHpYLoLVr14rf7zdvx+NxKSsrk+bmZoVdzR+RSEQASEdHh4iIDA0NSU5OjrS2tpo1ly9fFgASDAZFROTEiRNit9slFAqZNS0tLaJpmsRisfRuQBoNDw/L0qVLpb29XZ544gkzgLjP0sdSL8HGx8fR1dUFn89nrrPb7fD5fAgGgwo7mz8MwwDw9w/sdnV1YWJiImmfLVu2DF6v19xnwWAQK1asgNvtNmtqamoQjUbR29ubxu7Ty+/3o7a2NmnfANxn6WSpD6MODg4iHo8n/acDgNvtxpUrVxR1NX8kEgns2LED69evx/LlywEAoVAIDocDhYWFSbVutxuhUMismW6fTo1loqNHj+KTTz7B+fPn7xjjPksfSwUQ3Zvf78fFixfx17/+VXUr89rAwABeeukltLe3Y8GCBarbeahZ6iVYSUkJsrKy7piNCIfD8Hg8irqaHxoaGtDW1oaPP/4Y5eXl5nqPx4Px8XEMDQ0l1f/jPvN4PNPu06mxTNPV1YVIJILvfe97yM7ORnZ2Njo6OvDWW28hOzsbbreb+yxNLBVADocDVVVVCAQC5rpEIoFAIABd1xV2po6IoKGhAceOHcPJkyexZMmSpPGqqirk5OQk7bO+vj709/eb+0zXdfT09CASiZg17e3t0DQNlZWV6dmQNHrqqafQ09OD7u5uc1m9ejU2bdpk/pv7LE1UXwVP1dGjR8XpdMqRI0fk0qVL8uKLL0phYWHSbMTDZNu2beJyueTUqVNy/fp1c7l165ZZs3XrVvF6vXLy5Em5cOGC6Louuq6b41NTyhs2bJDu7m758MMP5Vvf+tZDNaX8j7NgItxn6WK5ABIR+d3vfider1ccDoesXbtWzp49q7olZQBMuxw+fNis+frrr+VXv/qVPPLII7Jw4UL56U9/KtevX096nC+//FKefvppyc3NlZKSEvn1r38tExMTad4adW4PIO6z9ODXcRCRMpa6BkREmYUBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKfN/h8fCMOw/fyQAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "200.0"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "play(True)[-1]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pt39]",
   "language": "python",
   "name": "conda-env-pt39-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
